Skip to content

Add byte-progress reporting for safetensors file loading - #427

Open
aleroot wants to merge 2 commits into
ml-explore:mainfrom
aleroot:loading_progress
Open

Add byte-progress reporting for safetensors file loading#427
aleroot wants to merge 2 commits into
ml-explore:mainfrom
aleroot:loading_progress

Conversation

@aleroot

@aleroot aleroot commented Jun 20, 2026

Copy link
Copy Markdown
Contributor

Proposed changes

Adds byte-level progress reporting for safetensors file loading so upper layers can drive precise model-loading progress UI.

This PR depends on ml-explore/mlx#3734.

The progress API in mlx-swift can report byte-level safetensors loading progress, but the truncated/failed-read correctness path needs the upstream MLX fix first. Without ml-explore/mlx#3734, CPU lazy-load read failures can be swallowed before they reach Swift.

Once ml-explore/mlx#3734 eventaully lands, I’ll update the MLX submodule reference here and rerun the mlx-swift test suite.

Checklist

Put an x in the boxes that apply.

  • I have read the CONTRIBUTING document
  • I have run pre-commit run --all-files to format my code / installed pre-commit prior to committing changes
  • I have added tests that prove my fix is effective or that my feature works
  • I have updated the necessary documentation (if needed)

@aleroot

aleroot commented Jun 25, 2026

Copy link
Copy Markdown
Contributor Author

This now depends on ml-explore/mlx#3742 , I will update this PR when that one is going to be merged.

Comment thread Source/MLX/IO.swift Outdated
Comment on lines +9 to +10
/// single load are delivered in monotonically increasing order.
public struct LoadProgress: Sendable, Equatable {

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Is it worth documenting that:

  • it may only read part of the file -- if safetensors are not consumed then it won't read them
  • I don't know if it is possible, but it could read part of the file twice and it could go over 100%

Not an issue but might be surprising (loading safetensors is lazy but not in an obvious way).

Comment thread Source/Cmlx/mlx-c

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I can't find the 00375d0 commit in mlx-c -- I presume this PR requires a newer mlx/mlx-c?

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Yes, before we can merge this, this other one need to be merged first, then this one needs to be updated.

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Today a new version of mlx has been released 🍾 , opened this one ml-explore/mlx-c#126

Introduces LoadProgress and new loadArrays(url:stream:progressHandler:) and
loadArraysAndMetadata(url:stream:progressHandler:) overloads that report
bytes read as lazy arrays are evaluated. Uses a custom mlx_io_reader vtable
backed by pread() so progress callbacks can be invoked from MLX worker
threads. Includes a unit test verifying monotonic progress from 0 to 1.
`loadWeights()` style helpers -- including the one in mlx-swift-lm -- call the
plain `loadArrays(url:)` / `loadArraysAndMetadata(url:)`, so a per-call
`progressHandler:` argument can never reach them without changing every caller
along the way.

Add `withLoadProgressHandler(_:_:)` (sync and async), a task local scoped
handler in the style of `withErrorHandler(_:_:)`. The plain file loading
functions report byte progress to it when one is installed, so an application
can drive a precise model loading progress bar around code it does not own:

    let container = try await withLoadProgressHandler({ tracker.update($0) }) {
        try await factory.loadContainer(from: directory, using: tokenizerLoader)
    }

`LoadProgress` gains the `url` of the file being read so progress can be
aggregated across the shards of a sharded model.

Also fix the SEEK_END case of the in-memory reader, which moved the offset
relative to the current position instead of the end of the data. mlx 0.32.1
seeks to the end of the stream to validate the tensor data offsets against the
size of the file, so `loadArrays(data:)` would fail there ("The JSON header is
N bytes long but the file is only 8 bytes").

Finally, restructure the truncated file test: a truncated file may now be
reported either eagerly, while the header is parsed, or lazily, when the arrays
are evaluated, and neither happens with the currently vendored mlx/mlx-c.
aleroot added a commit to aleroot/mlx-swift-lm that referenced this pull request Aug 18, 2026
Loading a model is dominated by reading the weights from disk, but nothing is
reported while that happens: `progressHandler` only covers the download, so an
application showing "loading..." has no way to draw an accurate progress bar,
and a large model can spend tens of seconds there.

Group the progress callbacks of a load in one value, `LoadProgressHandlers`:

    let container = try await LLMModelFactory.shared.loadContainer(
        from: directory, using: tokenizerLoader,
        progress: .weights { progress in
            print(progress.fractionCompleted)
        })

It is built on the scoped progress handler of mlx-swift, so `_load()` and the
model implementations are untouched -- the handler is installed around the load
and the plain `loadArrays(url:)` calls in `loadWeights()` report to it.
`ModelLoadProgressReporter` aggregates the byte progress that MLX reports per
file -- a model is frequently split into several shards, read concurrently --
into a single `Progress` for the whole model, and coalesces the updates, as MLX
reports roughly one per 4MB and the handler typically hops to the main actor.

Loading is lazy, so the weights are read while the model is evaluated at the end
of `loadWeights()`. The weights that `sanitize(weights:metadata:)` drops are
never evaluated, and therefore never read, so the aggregate can legitimately
stop short of the size of the files: completion is published once the load
returns.

The existing `progressHandler` parameter is unchanged and keeps reporting the
download; `progress.download` is called in addition to it when both are given.

Note: requires the scoped `withLoadProgressHandler(_:_:)` of
ml-explore/mlx-swift#427.
aleroot added a commit to aleroot/mlx-swift-lm that referenced this pull request Aug 18, 2026
Loading a model is dominated by reading the weights from disk, but nothing is
reported while that happens: `progressHandler` only covers the download, so an
application showing "loading..." has no way to draw an accurate progress bar,
and a large model can spend tens of seconds there.

Group the progress callbacks of a load in one value, `LoadProgressHandlers`:

    let container = try await LLMModelFactory.shared.loadContainer(
        from: directory, using: tokenizerLoader,
        progress: .weights { progress in
            print(progress.fractionCompleted)
        })

It is built on the scoped progress handler of mlx-swift, so `_load()` and the
model implementations are untouched -- the handler is installed around the load
and the plain `loadArrays(url:)` calls in `loadWeights()` report to it.
`ModelLoadProgressReporter` aggregates the byte progress that MLX reports per
file -- a model is frequently split into several shards, read concurrently --
into a single `Progress` for the whole model, and coalesces the updates, as MLX
reports roughly one per 4MB and the handler typically hops to the main actor.

Loading is lazy, so the weights are read while the model is evaluated at the end
of `loadWeights()`. The weights that `sanitize(weights:metadata:)` drops are
never evaluated, and therefore never read, so the aggregate can legitimately
stop short of the size of the files: completion is published once the load
returns.

The existing `progressHandler` parameter is unchanged and keeps reporting the
download; `progress.download` is called in addition to it when both are given.

Note: requires the scoped `withLoadProgressHandler(_:_:)` of
ml-explore/mlx-swift#427.
aleroot added a commit to aleroot/mlx-swift-lm that referenced this pull request Aug 18, 2026
Loading a model is dominated by reading the weights from disk, but nothing is
reported while that happens: `progressHandler` only covers the download, so an
application showing "loading..." has no way to draw an accurate progress bar,
and a large model can spend tens of seconds there.

Group the progress callbacks of a load in one value, `LoadProgressHandlers`:

    let container = try await LLMModelFactory.shared.loadContainer(
        from: directory, using: tokenizerLoader,
        progress: .weights { progress in
            print(progress.fractionCompleted)
        })

It is built on the scoped progress handler of mlx-swift, so `_load()` and the
model implementations are untouched -- the handler is installed around the load
and the plain `loadArrays(url:)` calls in `loadWeights()` report to it.
`ModelLoadProgressReporter` aggregates the byte progress that MLX reports per
file -- a model is frequently split into several shards, read concurrently --
into a single `Progress` for the whole model, and coalesces the updates, as MLX
reports roughly one per 4MB and the handler typically hops to the main actor.

Loading is lazy, so the weights are read while the model is evaluated at the end
of `loadWeights()`. The weights that `sanitize(weights:metadata:)` drops are
never evaluated, and therefore never read, so the aggregate can legitimately
stop short of the size of the files: completion is published once the load
returns.

The existing `progressHandler` parameter is unchanged and keeps reporting the
download; `progress.download` is called in addition to it when both are given.

Note: requires the scoped `withLoadProgressHandler(_:_:)` of
ml-explore/mlx-swift#427.
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants